%===============================================================================
% Run D4 Wavelet encoding and quantization through a range of threshold
% cutoff values and analyze the relative error energy with respect
% to the cutoff fraction and estimated compressed image size.
% 
% Arguments:
% numberTrials (int) - number of thresholds to analyze in range (default=10)
% cutoffMin (double) - lower threshold analysis range (default 0.75)
% cutoffMax (double) - lower threshold analysis range (default 0.99)
%===============================================================================
function analyzeD4Compression(numberTrials, cutoffMin, cutoffMax)
  if nargin<3
    cutoffMin = 0.75;
    cutoffMax = 0.99;
  end
  if nargin<1
    numberTrials=6;
  end
  
  pic = 'Pic.jpg';
  resolution = 256;

  relativeError = [];
  compressionRatio = [];
  cutoffRange = linspace(cutoffMin,cutoffMax,numberTrials); 

  i=1;
  for cutoff = cutoffRange
    msg = sprintf('%d: processing threshold cutoff: %f',i, cutoff);
    disp(msg);
    
    [relativeError(i), compressionRatio(i)] = ...
        D4CompressImage(pic, resolution, cutoff,i, numberTrials);

    msg = sprintf('    Relative Error: %f', relativeError(i));
    disp(msg);    
    msg = sprintf('    Compression Ratio: %f', compressionRatio(i));
    disp(msg);        
    msg = sprintf('    inv[Compression Ratio]: %f', 1/compressionRatio(i));
    disp(msg);            
    i = i+1;      
  end

  figure; % force the plot to appear
  
  subplot(2,1,1);  
  plot(compressionRatio, relativeError,'*');
  title('Compression Ratio Analysis');  
  xlabel('Compression Ratio (original size/compressed size)');
  ylabel('Relative Error Energy');

  subplot(2,1,2);
  plot(cutoffRange,relativeError,'*');
  xlim([cutoffMin cutoffMax]);
  title('Quantization Cutoff Threshold Analysis');
  xlabel('Quantization Cutoff Threshold Fraction');  
  ylabel('Relative Error Energy');  
  
%===============================================================================
% D4 Wavelet Image Compression
% Arguments: jpgFile (string)           default='Pic.jpg'
%            resolution (int)           default=256 i.e. 256x256 image
%            cutoff (double)            default=0.50 i.e. 0%      
%===============================================================================
function [relativeError, compressionRatio] = ...
    D4CompressImage(jpgFile, resolution, cutoff, count, numberTrials)
if nargin < 3
  cutoff = 0.50;
end
if nargin < 2
  resolution = 256;
end
if nargin < 1
  jpgFile = 'Pic.jpg';
end

numRows = resolution;
numCols = numRows;    

grayPic = getGrayImage(jpgFile, numRows, numCols);
grayPicEncoded = D4encode(grayPic);

quantStruct = struct([]); 
quantStruct = logQuantize(grayPicEncoded, cutoff);
grayPicEncodedQuantized = quantStruct.matrix;

% perform zig-zag scanning, Huffman encoding, and dump to file
[zipped, info] = homeBrewCompression(grayPicEncodedQuantized);

compressedPicSize = info.length * info.ratio;
originalPicSize = resolution^2;
compressionRatio = compressedPicSize/originalPicSize;

thresh = quantStruct.thresh;
lmaxt = quantStruct.lmaxt;

grayPicEncodedDequantized = logDequantize(grayPicEncodedQuantized, thresh, lmaxt);
grayPicDecompressed = D4decode(grayPicEncodedDequantized);

errPic = double(grayPic) - grayPicDecompressed;
errPicSqr = errPic .^ 2;
errPicEnergy = calcEnergy(errPic);
relativeError = relativeCodingError(grayPic, grayPicDecompressed);

% plot the decompressed images in a panel for comparison
subplot(ceil(numberTrials/2),2,count);
imagesc(grayPicDecompressed);
axis equal; 
xlim([1 256]);
ylim([1 256]);
colormap gray(256); 
msg = sprintf('Decompressed Image (compression ration = %f)', compressionRatio); 
title(msg);    

%===============================================================================
% Helper Functions
%===============================================================================

function relErr = relativeCodingError(refImage, codedImage)
  error = double(refImage) - double(codedImage);
  refEnergy = calcEnergy(refImage);
  errorEnergy = calcEnergy(error);
  relErr = errorEnergy/refEnergy;

function pic = getGrayImage(jpgFile, numRows, numCols)
  if nargin < 1
    jpgFile = 'Pic.jpg';
  end
  if nargin < 3
    numRows = 256; numCols = 256;  
  end

  picRGB = imread(jpgFile);

  imageDim = [numRows, numCols];

  % scale the RGB image - keep this around just for reference
  picRGBResized = imresize(picRGB, imageDim, 'bicubic', 'Antialiasing', true);
  % convert RGB image Gray scale
  pic = rgb2gray(picRGBResized);

%===============================================================================
function energy = calcEnergy(A)
  energy = sum(A(:) .^ 2);
%===============================================================================
function  dc = calcAverage(A)
  dc = sum(A(:)) / (size(A,2)^2);
%===============================================================================
function quantStruct = logQuantize(A, cutoff)
  % take all matrix entries, take abs, and arrange as a sorted vector X
  [m,n] = size(A);
  mn = m*n;  % number of entries in matrix
  X = reshape(A,mn,1);
  X = sort(abs(X));  
  maxA = max(X);
  if cutoff == 0
    thresh = 1;
    lmaxt = log2(maxA);    
  else
    thresh = X(ceil(cutoff*mn)); % threshold is matrix value at cutoff index
    lmaxt = log2(maxA/thresh);  
  end

  absA = abs(A); % take absolute value only once
  
  for i=1:m
    for j=1:n
      if absA(i,j) < thresh
        A(i,j) = 0;
      else
        % 256 level log quantization
        sign = A(i,j)/absA(i,j);
        ln = log2(absA(i,j)/thresh);
        q = ceil(127*ln/lmaxt);
        A(i,j) = sign*q;
      end
    end
  end
  
  quantStruct = struct('dummy',123); % hack - struct needs an initial field
  quantStruct.matrix = A;
  quantStruct.thresh = thresh;
  quantStruct.lmaxt = lmaxt;
%===============================================================================  
function matrix = logDequantize(A, thresh, lmaxt)
  [m,n] = size(A);
  absA = abs(A);
  for i=1:m
    for j=1:n
      if absA(i,j) >0
        sign = A(i,j)/absA(i,j);
        fp = absA(i,j)*lmaxt/127;
        A(i,j) = sign * thresh * 2^fp;
      end
    end
  end
  matrix = A;
%===============================================================================  
% D4 wavelet transform encoding
% Compute the D4 coefficients of the input matrix 
% The matrix is assumed to be square
%===============================================================================
function matrix = D4encode(A)

matrixSize = size(A);
length = matrixSize(1);
log2length = ceil(log2(length));
iterations = log2length-1;

B = double(A); % initialize the coefficient matrix

% recursively iterate to build output matrix B
% halve the length of each dimension in each iteration
for j=1:iterations
  k = log2length+1-j;
  length = 2^k;
  C = B(1:length, 1:length);
  
  H = D4transform(length);

  % construct Permutation matrix instance
  I = eye(length);   
  PT = I([1:2:length],:);  % averages
  PB = I([2:2:length],:);  % differences
  P = [PT; PB];

  % apply transform filter
  B(1:length,1:length) = P*H*C*H'*P';

  % plot octave stages of tranform for debugging
  %{   
  subplot(2,ceil(iterations/2),j);
  imagesc(B); 
  axis equal; 
  msg = sprintf('D4 encoding iteration: %d',j);
  title(msg);
  colormap gray(256); 
  %}   
end
matrix = B;
%===============================================================================
% Inverse D4 transform
% Given the D4 coefficients, reconstitute the matrix
%===============================================================================
function matrix = D4decode(A)

matrixSize = size(A);
length = matrixSize(1);
log2length = ceil(log2(length));
iterations = log2length-1;

% recursively iterate to build output matrix B
% double the length of each dimension in each iteration

C = double(A(1,1));
%for j=1:iterations
for j=2:iterations+1
  length = 2^j;
  halfLength = 2^(j-1);
  B = double(A(1:length, 1:length));

  B(1:halfLength, 1:halfLength) = C;
  
  H = D4transform(length);
  
  % construct Permutation matrix instance
  I = eye(length);   
  PT = I([1:2:length],:);  % averages
  PB = I([2:2:length],:);  % differences
  P = [PT; PB];

  % apply transform filter
  C = H'*P'*double(B)*P*H;
end
matrix = C;

%===============================================================================
function H = D4transform(length)
% construct master D4 transform filter matrix
% we will extract smaller Haar sub-matrices from this

h0 = (1+sqrt(3))/(4*sqrt(2));
h1 = (3+sqrt(3))/(4*sqrt(2));
h2 = (3-sqrt(3))/(4*sqrt(2));
h3 = (1-sqrt(3))/(4*sqrt(2));

% Quadrature Mirror Filter (QMF) Bank
h = [h0 h1 h2 h3];    % low pass filter  LPF
g = [h3 -h2 h1 -h0];  % high pass filter HPF

hg = [h;g];

H = zeros(length);
for i=1:ceil(length/2)
  index = (i-1)*2+1;
  if index+3 <= length
    H(index:index+1,index:index+3) = hg;        
  else
    % wrap-around coefficients accounts for filter periodicity
    H(index:index+1,index:index+1) = hg(1:2,1:2);    
    H(index:index+1,1:2) = hg(1:2,3:4);        
  end
end

%===============================================================================
% count the number of non-zero array elements
function count = compress(A)
  [m,n] = size(A);
  count = 0;
  for i=1:m
    for j=1:n
      if A(i,j) ~= 0
        count = count + 1;
      end
    end
  end

  
